# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import enum
import functools
import typing
import uuid

import requests
import tornado.gen

import coin.flow.topic_map as topic_map
import coin.flow.subscriber as flow_subscriber
import coin.proto.coin_market_enums_pb2 as coin_enum
import coin.proto.coin_order_enums_pb2 as coin_order

from coin.base.config import Config
from coin.base.timestamp import get_timestamp
from coin.exchange.base.order_gateway_ready import PositionOrderGatewayReady
from coin.exchange.base.rest.task_scheduler import Scheduler, wrap_async_task
from coin.exchange.kraken.kr_rest.currency import KrakenCurrency
from coin.exchange.kraken_futures.kr_rest.futures_product import KrakenFuturesProduct
from coin.exchange.kraken_futures.kr_rest.native_private_client import KrakenFuturesNativePrivateClient
from coin.exchange.kraken_futures.kr_rest.private_client import KrakenFuturesPrivateParser
from coin.exchange.util.rest_rate_limit import RestRateLimiter
from coin.proto.coin_order_gateway_pb2 import OrderGatewayConfig, OrderEvent
from coin.proto.coin_query_pb2 import ProductPosition, ProductFillElement

from coin.exchange.kraken_futures.ws.ws_private_client import (
    KrakenFuturesWsPrivateClient,
    ChannelType,
)
from coin.exchange.kraken_futures.kr_rest.native_private_client import (
    KrakenFuturesOrderSide,
    KrakenFuturesOrderType,
    KrakenFuturesOrderStatus,
    KrakenFuturesFillType,
    KrakenFuturesErrorType,
    KrakenFuturesSubmitStatus,
    KrakenFuturesCancelStatus,
    KrakenFuturesWsOrderUpdateReason,
)
from coin.exchange.base.rest_order_gateway_base import (
    FillManager,
    Order,
    OrderType,
    OrderSide,
    OrderGatewayStatus,
    RestOrderGateway,
    BasePositionOrderGatewayInfo,
    create_order_internal,
    exchange_message_handler,
    pretty_json,
    PrivateExchangeMessage,
    to_proto_order_side,
)


def create_order_from_open_order_message(message, timestamp, order_id=None):
  """
  {
    "receivedTime": "2016-02-25T09:45:53.601Z",
    "status": "partiallyFilled",
    "order_id": "c18f0c17-9971-40e6-8e5b-10df05d422f0",
    "cliOrdId": "dd84e375-474a-48df-bc54-a13b02ef6e8e",
    "orderType": "stp",
    "symbol": "fi_xbtusd_180615",
    "side": "buy",
    "unfilledSize": 5000,
    "filledSize": 5000,
    "limitPrice": 4255,
    "stopPrice": 4202
  }
  """
  product = KrakenFuturesProduct.FromStrNativeProduct(message['symbol'])
  price_orig = float(message['limitPrice'])
  filled_qty = float(message['filledSize'])
  unfilled_qty = float(message['unfilledSize'])
  qty_orig = filled_qty + unfilled_qty
  exchange_order_id = message['order_id']
  qty_accum_fill = filled_qty
  price_avg_fill = price_orig  # Since filled price info not provided
  status = KrakenFuturesOrderStatus(message['status'])
  order_side = getattr(OrderSide, KrakenFuturesOrderSide(message['side']).name)
  order_type = OrderType.LIMIT

  fill = FillManager(
      price_orig=price_orig,
      qty_orig=qty_orig,
      price_last_fill=price_avg_fill,
      qty_last_fill=qty_accum_fill,
      price_avg_fill=price_avg_fill,
      qty_accum_fill=qty_accum_fill,
      timestamp_last_fill=timestamp,
  )
  order = Order(order_id=order_id,
                product=product,
                price=price_orig,
                qty=qty_orig,
                order_side=order_side,
                order_type=order_type,
                tag=None,
                accepted=True,
                cancel_sent=False,
                internal=create_order_internal(
                    timestamp=timestamp,
                    exchange_order_id=exchange_order_id,
                    fill=fill,
                    status=status,
                ))
  return order


def create_order_info_from_ws_message(message, timestamp, order_id=None):
  """
  {
    "instrument": "FI_XBTUSD_180316",
    "time": 1516888562000,
    "qty": 480.0,
    "filled": 0.0,
    "limit_price": 12091.0,
    "stop_price": 0.0,
    "type": "limit",
    "order_id": "97325f00-b901-46cb-9029-6b52209e3473",
    "direction": 1
  }
  """
  order_sides = [OrderSide.BUY, OrderSide.SELL]
  product = KrakenFuturesProduct.FromStrNativeProduct(message['instrument'])
  price_orig = float(message['limit_price'])
  qty_orig = float(message['qty'])
  exchange_order_id = message['order_id']
  qty_accum_fill = float(message['filled'])
  price_avg_fill = price_orig  # Since filled price info not provided
  status = None
  if 'status' in message:
    status = KrakenFuturesOrderStatus(message['status'])
  order_side = order_sides[message['direction']]
  assert message['type'] == 'limit', message
  order_type = OrderType.LIMIT

  fill = FillManager(
      price_orig=price_orig,
      qty_orig=qty_orig,
      price_last_fill=price_avg_fill,
      qty_last_fill=qty_accum_fill,
      price_avg_fill=price_avg_fill,
      qty_accum_fill=qty_accum_fill,
      timestamp_last_fill=timestamp,
  )
  order = Order(order_id=order_id,
                product=product,
                price=price_orig,
                qty=qty_orig,
                order_side=order_side,
                order_type=order_type,
                tag=None,
                accepted=True,
                cancel_sent=False,
                internal=create_order_internal(
                    timestamp=timestamp,
                    exchange_order_id=exchange_order_id,
                    fill=fill,
                    status=status,
                ))
  return order


class KrakenFuturesOrderGatewayInfo(BasePositionOrderGatewayInfo):
  CurrencyType = KrakenCurrency
  ProductType = KrakenFuturesProduct
  BalanceType = float
  PositionType = float

  def __init__(self,
               currencies: typing.List[KrakenCurrency],
               products: typing.List[KrakenFuturesProduct]):
    super().__init__(currencies, products)
    self._unrealized_pnl = None

  def update_balance(self, message):
    msg = KrakenFuturesPrivateParser.parse_rest_trade_balance(message)
    for balance in msg.each_balance:
      currency = KrakenCurrency.FromStrNativeCurrency(balance.currency_native)
      if currency in self.currencies:
        self.set_balance(currency, balance.total, balance)
    self._unrealized_pnl = KrakenFuturesPrivateParser.parse_unrealized_pnl(message)

  def update_position(self, message):
    for product in self.products:
      position = ProductPosition(symbol=product.symbol,
                                 long_position=0,
                                 short_position=0,
                                 net_position=0)
      self.set_position(product, 0.0, position)

    msg = KrakenFuturesPrivateParser.parse_rest_account_position(message)
    for position in msg.each_position:
      product = KrakenFuturesProduct.FromStrProduct(position.symbol)
      if product in self.products:
        self.set_position(product, position.net_position, position)

  # Margin balance includes unrealized pnl, but this is provided by exchange
  # and could be out of date or slow.
  # DO NOT USE THIS IN TRADING, ONLY ACCEPTABLE IN PNL DISPLAY.
  def get_margin_balance(self, currency: str):
    self.get_balance(currency) + self._unrealized_pnl[currency]


class KrakenFuturesOrderGateway(RestOrderGateway):
  _exchange = 'Kraken'
  _market_type = 'Futures'
  _api_version = None

  OGInfoType = KrakenFuturesOrderGatewayInfo

  class _MsgType(enum.IntEnum):
    REST_ACCOUNT_BALANCE = 1
    REST_ACCOUNT_POSITION = 2
    REST_OPEN_ORDERS = 3
    REST_FILLS = 4
    WS_ACCOUNT_POSITION = 5
    WS_OPEN_ORDERS = 6
    WS_FILLS = 7
    WS_HEARTBEAT = 8
    WS_OPEN_ORDERS_SNAPSHOT = 9

  def __init__(self,
               currencies: typing.List[KrakenCurrency],
               products: typing.List[KrakenFuturesProduct],
               config: Config,
               og_config=None,
               logger=None):
    super().__init__(currencies, products, config, og_config, logger)
    self._og_ready = PositionOrderGatewayReady()
    self._og_ready.set_status(OrderGatewayStatus.INITIALIZING)

    # Following instance attributes are initialized in start method.
    self._flow_sub = None
    self._ioloop = None
    self._publisher = None
    self._scheduler = None
    self._private_client = None
    self._ws_private_client = None
    self._rest_rate_limiter = RestRateLimiter(delay=10)

    self._og_config = OrderGatewayConfig(
        account_query_period=10,
        user_query_period=2,
        orders_snapshot_query_period=10,
        fill_query_period=2,
        use_fill_checker=True,
    )
    if og_config is not None:
      self._og_config.MergeFrom(og_config)
    self._og_logger.exchange_type = coin_enum.Kraken

  def submittable(self):
    return self.is_ready() and self._rest_rate_limiter.is_available()

  def start(self, queue, ioloop, flow_sub=None):
    api_host = self._og_config.rest_root_url if \
               self._og_config.HasField('rest_root_url') else None
    disable_ssl = self._og_config.rest_disable_ssl

    topic_id = topic_map.search_topic_id('kraken_futures_rest_private_1')[0][1]
    self._ioloop = ioloop

    self._publisher = queue.get_writer(topic_id=topic_id)
    self._scheduler = Scheduler(ioloop=ioloop, logger=self._logger)

    self._private_client = KrakenFuturesNativePrivateClient(key_file=self._config.key_file,
                                                            api_host=api_host,
                                                            disable_ssl=disable_ssl,
                                                            use_async=True)
    self._prepare_ws_private_client(start=True)
    self._install_periodic_query()

    # Subscribe to flow for feed.
    self._flow_sub = flow_sub or flow_subscriber.from_queue(queue)
    self._flow_sub.subscribe('kraken_futures_rest_private_1', self.on_private_exchange_msg)

  @tornado.gen.coroutine
  def _install_periodic_query(self):
    assert self._scheduler is not None, 'self._scheduler is not set!'
    assert self._publisher is not None, 'self._publisher is not set!'

    # query balance
    query_balance_task = wrap_async_task(
        self._private_client.query_accounts,
        msg_type=self._MsgType.REST_ACCOUNT_BALANCE,
        publisher=self._publisher,
    )
    self._scheduler.add_periodic_task(
        period=self._og_config.account_query_period,
        job_func=query_balance_task,
    )

    yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
    # query positions
    query_position_task = wrap_async_task(
        self._private_client.query_open_position,
        msg_type=self._MsgType.REST_ACCOUNT_POSITION,
        publisher=self._publisher,
    )
    self._scheduler.add_periodic_task(
        period=self._og_config.user_query_period,
        job_func=query_position_task,
    )

    yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
    # query open orders
    query_orders_task = wrap_async_task(
        self._private_client.query_open_orders,
        msg_type=self._MsgType.REST_OPEN_ORDERS,
        publisher=self._publisher,
    )
    self._scheduler.add_periodic_task(period=self._og_config.orders_snapshot_query_period,
                                      job_func=query_orders_task)

    yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
    # query fills
    query_orders_info_task = wrap_async_task(
        self._private_client.query_fills,
        msg_type=self._MsgType.REST_FILLS,
        publisher=self._publisher,
    )
    self._scheduler.add_periodic_task(
        period=self._og_config.fill_query_period,
        job_func=query_orders_info_task,
    )
    self.start_scheduler()

  def _prepare_ws_private_client(self, start=True):
    if self._ws_private_client is not None:
      self._ws_private_client.stop()

    close_callback = functools.partial(self._prepare_ws_private_client, True)
    self._ws_private_client = KrakenFuturesWsPrivateClient(
        self._config.key_file,
        ioloop=self._ioloop,
        notification_callback=self._notification_callback,
        close_callback=close_callback,
        on_ready=self._on_ws_private_client_ready,
    )
    if start:
      self._ws_private_client.start()

  def _notification_callback(self, ws_msg):
    timestamp = get_timestamp()
    feed_type = ChannelType(ws_msg['feed'])
    if feed_type == ChannelType.OPEN_ORDERS:
      msg_type = self._MsgType.WS_OPEN_ORDERS
    elif feed_type == ChannelType.OPEN_ORDERS_SNAPSHOT:
      msg_type = self._MsgType.WS_OPEN_ORDERS_SNAPSHOT
    elif feed_type == ChannelType.FILLS:
      msg_type = self._MsgType.WS_FILLS
    elif feed_type == ChannelType.FILLS_SNAPSHOT:
      self._logger.info('Snap shot fill received!')
      return
    elif feed_type == ChannelType.HEARTBEAT:
      msg_type = self._MsgType.WS_HEARTBEAT
    else:
      self._logger.error('Unhandled message type: %s', ws_msg)
      return

    exchange_message = PrivateExchangeMessage(msg_type=msg_type,
                                              data=ws_msg,
                                              received_timestamp=timestamp,
                                              request_params=None,
                                              request_timestamp=None)
    self._publisher.write(timestamp, exchange_message)

  @tornado.gen.coroutine
  def _on_ws_private_client_ready(self):
    channels = [
        ChannelType.HEARTBEAT,
        ChannelType.OPEN_ORDERS,
        ChannelType.FILLS,
    ]
    for channel in channels:
      yield self._ws_private_client.subscribe(channel)

  def _async_submit_order(self, order, post_only=None):
    formatter = self._order_formatter[order.product]
    if post_only:
      order_type = KrakenFuturesOrderType.POST_ONLY
    else:
      order_type = KrakenFuturesOrderType.LIMIT

    client_order_id = str(uuid.uuid1())
    params = {
        'order_type': order_type,
        'symbol': order.product.native_symbol,
        'limit_price': formatter.format_price(order.price),
        'order_side': getattr(KrakenFuturesOrderSide, order.order_side.name),
        'order_size': formatter.format_qty(order.qty),
        'client_order_id': client_order_id,
    }
    fut = self._private_client.place_order(**params)
    return fut

  def _async_cancel_order(self, order):
    assert order.internal.exchange_order_id is not None, order
    fut = self._private_client.cancel_order(order_id=order.internal.exchange_order_id)
    return fut

  def _handle_submit_response(self, order, response):
    """
    {
      "result": "success",
      "sendStatus": {
        "order_id": "8f7d9907-0e43-444f-a39c-5ba59df69bf6",
        "status": "placed",
        "receivedTime": "2019-03-20T05:21:45.259Z"
      },
      "serverTime": "2019-03-20T05:21:45.260Z"
    }
    {
      "result": "success",
      "sendStatus": {
        "status": "invalidSize"
       },
       "serverTime": "2019-03-20T05:20:01.791Z"
    }
    """
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      self._order_manager.remove_order(order, ignore_error=True)
      return OrderEvent.ORDER_REJECTED

    message = response.json()
    if message['result'] != 'success':
      self._handle_error_message(message)
      self._order_manager.remove_order(order, ignore_error=True)
      return OrderEvent.ORDER_ERROR

    status_info = message['sendStatus']
    try:
      status = KrakenFuturesSubmitStatus(status_info['status'])
    except ValueError:
      self._logger.error('Unknown status value: %s', message)
      status = KrakenFuturesSubmitStatus.UNKNOWN
    except KeyError:
      self._logger.error('Weird! Submit status field is missing! %s', message)
      status = KrakenFuturesSubmitStatus.UNKNOWN

    if status not in (KrakenFuturesSubmitStatus.PLACED,
                      KrakenFuturesSubmitStatus.ATTEMPTED,
                      KrakenFuturesSubmitStatus.UNKNOWN):
      self._order_manager.remove_order(order)
      return OrderEvent.ORDER_REJECTED

    exchange_order_id = status_info.get('order_id')
    if exchange_order_id is None:
      self._logger.error('Weird! Accepted order without order id! %s', message)
      self._order_manager.remove_order(order, ignore_error=True)
      return OrderEvent.ORDER_REJECTED

    prev_order = self._order_manager.get(exchange_order_id=exchange_order_id)
    if prev_order is None:
      self._order_manager.set_order(order, exchange_order_id=exchange_order_id)
    else:
      if prev_order.order_id != order.order_id:
        self._logger.info('order id replaced: %s => %s', prev_order.order_id, order.order_id)
        self._order_manager.remove_order(prev_order)
        self._order_manager.set_order(order, exchange_order_id=exchange_order_id)
        self._order_manager.remove_order(order)
        prev_order.order_id = order.order_id
        self._order_manager.add_order(prev_order)
    return OrderEvent.ORDER_ACCEPTED

  def _handle_cancel_response(self, order, response):
    """
    {
      "result": "success",
      "cancelStatus": {
        "status": "notFound/filled",
        "receivedTime": "2019-03-20T04:19:51.777Z"
      },
      "serverTime": "2019-03-20T04:19:51.777Z"
    }
    {
      "result": "success",
      "cancelStatus": {
        "status": "cancelled",
        "order_id": "8f7d9907-0e43-444f-a39c-5ba59df69bf6",
        "receivedTime": "2019-03-20T06:24:30.755Z"
      },
      "serverTime": "2019-03-20T06:24:30.755Z"
    }
    """
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return OrderEvent.CANCEL_ERROR

    message = response.json()
    if message['result'] != 'success':
      self._handle_error_message(message)
      return OrderEvent.CANCEL_ERROR

    status_info = message['cancelStatus']
    try:
      status = KrakenFuturesCancelStatus(status_info['status'])
    except ValueError:
      self._logger.error('Unknown status value: %s', message)
      status = KrakenFuturesCancelStatus.UNKNOWN
    except KeyError:
      self._logger.error('Weird! Cancel status field is missing! %s', message)
      status = KrakenFuturesCancelStatus.UNKNOWN

    exchange_order_id = status_info.get('order_id')
    if exchange_order_id is None and status not in (KrakenFuturesCancelStatus.NOT_FOUND,
                                                    KrakenFuturesCancelStatus.FILLED):
      self._logger.error('Weird! Cancel accepted order without order id! %s', message)
      order.cancel_sent = False
      return OrderEvent.CANCEL_ERROR
    elif exchange_order_id != order.internal.exchange_order_id:
      if status in (KrakenFuturesCancelStatus.NOT_FOUND, KrakenFuturesCancelStatus.FILLED):
        self._logger.info('Canceling NOT_FOUND or FILLED order, order status: %s', status)
      else:
        order.cancel_sent = False
        self._logger.error('Insane response, exchange order id not consistent: %s, %s',
                           order,
                           message)
        return OrderEvent.CANCEL_ERROR

    if status == KrakenFuturesCancelStatus.NOT_FOUND:
      self._logger.error('Cannot found order %s(%s)',
                         order.order_id,
                         order.internal.exchange_order_id)
      if order in self._order_manager:
        if self._order_manager.is_dead_order(order):
          self._logger.warning('Order (%s) is dead already!', order.order_id)
        else:
          self._order_manager.set_order(order, cancel_confirmed=True)
      else:
        self._logger.error('Order not found in order manager! Unexpected! %s', pretty_json(message))
      return OrderEvent.CANCEL_CONFIRMED
    elif status == KrakenFuturesCancelStatus.CANCELED:
      if order in self._order_manager:
        if not self._order_manager.is_dead_order(order):
          self._order_manager.set_order(order, cancel_confirmed=True)
      else:
        self._logger.error('Order not found in order manager! Unexpected! %s', pretty_json(message))
      return OrderEvent.CANCEL_CONFIRMED
    elif status == KrakenFuturesCancelStatus.FILLED:
      self._logger.error('Cancel failed since order is filled already! %s, %s',
                         order.order_id,
                         order.internal.order_id)
      return OrderEvent.CANCEL_REJECTED
    else:
      self._logger.error('Unknown order status! %s, %s',
                         order.order_id,
                         order.internal.exchange_order_id)

  def _update_order(self, order_info):
    exchange_order_id = order_info.internal.exchange_order_id
    order = self._order_manager.get(exchange_order_id=exchange_order_id)
    if order is None:
      order_info.order_id = self._order_id_generator.gen_order_id(is_foreign=True)
      self._order_manager.add_order(order_info)

  def _handle_error_response(self, response):
    if response.status_code == 504:
      self._logger.error('504 Gateway Timeout error!')
    elif response.status_code == 429:
      self._rest_rate_limiter.set_error(response.status_code)
      self._logger.error('Rate limit hit (429)!')

    self._logger.error('response error! %s, %s, %s',
                       response.status_code,
                       response.content,
                       response.url)

  def _handle_error_message(self, message):
    self._logger.error('result is not success! %s', message)
    error = KrakenFuturesErrorType(message['result'])
    self._rest_rate_limiter.set_error(error)

  def _handle_single_fill(self, fill_info, timestamp):
    # noinspection SpellCheckingInspection
    """
    {
      "fill_id": "742ce7d6-65a2-4323-90b2-dca9067f68d2",
      "symbol": "pv_xrpxbt",
      "side": "buy",
      "order_id": "faf61cb0-0b28-4e6a-b5e5-24c77bebfd52",
      "size": 1,
      "price": 8.049e-05,
      "fillTime": "2019-03-15T11:07:45.000Z",
      "fillType": "maker"
    }
    """
    fill_id = fill_info['fill_id']
    if fill_id in self._processed_fills:
      return

    native_symbol = fill_info.get('symbol')
    if native_symbol is not None:
      product = KrakenFuturesProduct.FromStr(fill_info['symbol'])
    else:
      self._logger.error('Unexpected fill info %s\n', pretty_json(fill_info))
      return

    if product not in self._og_info.products:
      return

    exchange_order_id = fill_info['order_id']
    order = self._order_manager.get(exchange_order_id=exchange_order_id)
    if order is None:
      self._logger.error('Order not found!, fill_info=%s', fill_info)
      return

    fill_type = KrakenFuturesFillType(fill_info['fillType'])
    if fill_type not in (KrakenFuturesFillType.MAKER_FILL_TYPE,
                         KrakenFuturesFillType.TAKER_FILL_TYPE):
      self._logger.warning('fill_type: %s', fill_type)
      fill_type = None
    else:
      fill_type = getattr(coin_order, fill_type.name)

    price_last_fill = float(fill_info['price'])
    qty_last_fill = float(fill_info['size'])
    last_fill = order.internal.fill.update_by_last_fill(price_last_fill, qty_last_fill)
    if last_fill:
      fill_proto = ProductFillElement(fill_id=fill_info['fill_id'],
                                      price=last_fill.price_last_fill,
                                      qty=last_fill.qty_last_fill,
                                      order_id=order.order_id,
                                      side=to_proto_order_side(order.order_side),
                                      exchange_order_id=exchange_order_id,
                                      fill_type=fill_type,
                                      symbol=order.product.symbol)
      self.process_fill(fill_proto, timestamp)
    else:
      self._logger.error('Filled qty more than the orig_qty: %s', order.order_id)
    self._processed_fills.add(fill_id)

  @exchange_message_handler(_MsgType.REST_OPEN_ORDERS)
  def _handle_rest_open_orders(self, exchange_message: PrivateExchangeMessage):
    timestamp = get_timestamp()
    response = exchange_message.data
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    message = response.json()
    if message['result'] != 'success':
      self._handle_error_message(message)
      return

    exchange_open_order_dict = {}
    for entry in message['openOrders']:
      native_symbol = entry['symbol']
      product = KrakenFuturesProduct.FromStrNativeProduct(native_symbol)
      if product not in self._og_info.products:
        continue  # ignore

      order = create_order_from_open_order_message(entry, timestamp)
      exchange_order_id = order.internal.exchange_order_id
      prev_order = self._order_manager.get(exchange_order_id=exchange_order_id)
      if prev_order is None:
        order_id = self._order_id_generator.gen_order_id(is_foreign=True)
        order.order_id = order_id
      else:
        order.order_id = prev_order.order_id
      exchange_open_order_dict[exchange_order_id] = order
    self._order_manager.copy_from_exchange_open_orders(exchange_order_dict=exchange_open_order_dict,
                                                       products=self._og_info.products)

    if not self.is_ready():
      self._og_ready.set_orders_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.REST_ACCOUNT_BALANCE)
  def _handle_rest_balance(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    message = response.json()
    if message['result'] != 'success':
      self._handle_error_message(message)
      return

    self._og_info.update_balance(message)
    self._og_logger.write_balance(self._og_info.balance_as_proto())
    if not self.is_ready():
      self._og_ready.set_balance_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.REST_ACCOUNT_POSITION)
  def _handle_rest_position(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    message = response.json()
    if message['result'] != 'success':
      self._handle_error_message(message)
      return

    self._og_info.update_position(message)
    self._og_logger.write_position(self._og_info.position_as_proto())
    if not self.is_ready():
      self._og_ready.set_position_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.REST_FILLS)
  def _handle_rest_fills(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    timestamp = exchange_message.received_timestamp
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    message = response.json()
    if message['result'] != 'success':
      self._handle_error_message(message)
      return

    is_first_query = not bool(self._processed_fills)
    if is_first_query:
      self._processed_fills.add(None)
      for fill_info_msg in message['fills']:
        fill_id = fill_info_msg['fill_id']
        self._processed_fills.add(fill_id)
      return

    for fill_info_msg in message['fills']:
      self._handle_single_fill(fill_info_msg, timestamp)

  @exchange_message_handler(_MsgType.WS_OPEN_ORDERS_SNAPSHOT)
  def _handle_ws_open_orders_snapshot(self, exchange_message: PrivateExchangeMessage):
    timestamp = exchange_message.received_timestamp
    data = exchange_message.data
    for order_msg in data['orders']:
      order_info = create_order_info_from_ws_message(order_msg, timestamp)
      if order_info in self._og_info.products:
        self._update_order(order_info)

  @exchange_message_handler(_MsgType.WS_OPEN_ORDERS)
  def _handle_ws_open_orders(self, exchange_message: PrivateExchangeMessage):
    timestamp = exchange_message.received_timestamp
    data = exchange_message.data
    reason = KrakenFuturesWsOrderUpdateReason(data.get('reason', 'unknown'))
    if reason == KrakenFuturesWsOrderUpdateReason.NEW_PLACED_ORDER_BY_USER:
      order_msg = data['order']
      order_info = create_order_info_from_ws_message(order_msg, timestamp)
      if order_info.product not in self._og_info.products:
        return
      self._update_order(order_info)
      exchange_order_id = order_msg['order_id']
    else:
      exchange_order_id = data['order_id']

    order = self._order_manager.get(exchange_order_id=exchange_order_id)
    if order is None:
      self._logger.error('Order not found in order manager! %s', pretty_json(data))
      return

    if reason == KrakenFuturesWsOrderUpdateReason.CANCELLED_BY_USER:
      self._order_manager.set_order(order, cancel_confirmed=True)
    elif reason == KrakenFuturesWsOrderUpdateReason.FULL_FILL:
      self._order_manager.set_order(order, fully_filled=True)

  @exchange_message_handler(_MsgType.WS_FILLS)
  def _handle_ws_fills(self, exchange_message: PrivateExchangeMessage):
    timestamp = exchange_message.received_timestamp
    data = exchange_message.data
    for fill_info in data['fills']:
      fill_info['size'] = fill_info['qty']
      fill_info['fillType'] = fill_info['fill_type']
      fill_info['symbol'] = fill_info['instrument']
      self._handle_single_fill(fill_info, timestamp)

  @exchange_message_handler(_MsgType.WS_HEARTBEAT)
  def _handle_heartbeat(self, exchange_message: PrivateExchangeMessage):
    data = exchange_message.data
    self._logger.debug('heartbeat: %s', data)
